import akshare as ak
import pandas as pd
from sqlalchemy import create_engine, text
import time

# ================= MySQL 连接配置 =================
engine = create_engine("mysql+pymysql://root:123456@localhost:3306/cn_stock_data")

# ================= 读取股票代码 =================
with engine.connect() as conn:
    result = conn.execute(text("SELECT stock_code FROM company_info"))
    codes = [row[0] for row in result.fetchall()]

# 补全为6位股票代码
codes = [code.zfill(6) for code in codes]

print(f"一共获取到 {len(codes)} 只股票代码")

# ================= 下载并写入 =================
for symbol in codes:
    try:
        print(f"正在下载 {symbol} 的数据...")

        stock_df = ak.stock_zh_a_hist(
            symbol=symbol,
            period="daily",
            start_date="20250301",
            end_date="20250828",
            adjust="hfq"  # 前复权，可按需修改
        )

        if stock_df.empty:
            print(f"⚠️ {symbol} 没有数据，跳过")
            time.sleep(2)
            continue

        # 删除可能存在的 index 列
        if "Unnamed: 0" in stock_df.columns:
            stock_df = stock_df.drop(columns=["Unnamed: 0"])

        # 列名映射到数据库表字段
        stock_df = stock_df.rename(columns={
            "日期": "trade_date",
            "股票代码": "stock_code",
            "开盘": "open_price",
            "收盘": "close_price",
            "最高": "high_price",
            "最低": "low_price",
            "成交量": "volume",
            "成交额": "amount",
            "振幅": "amplitude",
            "涨跌幅": "change_rate",
            "涨跌额": "change_amount",
            "换手率": "turnover_rate"
        })

        # 转换日期格式
        stock_df["trade_date"] = pd.to_datetime(stock_df["trade_date"]).dt.date

        # 写入 MySQL
        stock_df.to_sql(
            "his_daily_market_data",
            con=engine,
            if_exists="append",  # 追加
            index=False
        )

        print(f"✅ {symbol} 写入完成，{len(stock_df)} 条记录")

    except Exception as e:
        print(f"❌ 下载 {symbol} 出错: {e}")

    # 每只股票下载完暂停 2 秒
    time.sleep(2)
